package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"strings"

	"github.com/Shopify/sarama"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

func main() {
	path := flag.String("path", "/Users/liruopeng/Downloads/suricata-all", "pcap path")
	broker := flag.String("broker", "10.36.8.128:9092", "kafka broker list")
	flag.Parse()
	var fileList []string

	if *path == "" {
		fmt.Println("Please input the pcap file path.")
		os.Exit(0)
	}
	if *broker == "" {
		fmt.Println("Please input the Kafka broker url.")
		os.Exit(0)
	}

	if err := findPcap(*path, &fileList); err != nil {
		panic(err)
	}
	fmt.Println(fileList)
	producer := initProducer(*broker)

	// Trap SIGINT to trigger a shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)

	var (
		successes, errors int
	)

	go func() {
		for range (*producer).Successes() {
			successes++
		}
	}()

	go func() {
		for err := range (*producer).Errors() {
			log.Println(err)
			errors++
		}
	}()

	for _, path := range fileList {
		sendToKafka(path, producer)
	}
WaitLoop:
	for {
		select {
		case <-signals:
			(*producer).AsyncClose()
			break WaitLoop
		}
	}
	fmt.Println(successes, errors)
}

func findPcap(path string, fileList *[]string) error {
	err := filepath.Walk(path, func(path string, f os.FileInfo, err error) error {
		if !f.IsDir() && strings.HasSuffix(f.Name(), ".pcap") {
			*fileList = append(*fileList, path)
		}
		return nil
	})
	return err
}

func initProducer(url string) *sarama.AsyncProducer {
	config := sarama.NewConfig()
	config.Producer.MaxMessageBytes = 1000000
	// config.Producer.Return.Successes = true
	producer, err := sarama.NewAsyncProducer([]string{url}, config)
	if err != nil {
		panic(err)
	}
	return &producer
}

func sendToKafka(path string, producer *sarama.AsyncProducer) {
	handle, err := pcap.OpenOffline(path)
	if err != nil {
		log.Fatal(err)
	}
	defer handle.Close()
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())

	for packet := range packetSource.Packets() {
		message := &sarama.ProducerMessage{Topic: "pcapAttack", Value: sarama.StringEncoder(getPacketInfo(packet))}

		(*producer).Input() <- message

	}
}

func getPacketInfo(packet gopacket.Packet) string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "%d|", len(packet.Data()))
	fmt.Fprintf(&buf, "%d|", packet.Metadata().Timestamp.Unix())

	// 网络接口层
	ethernetLayer := packet.Layer(layers.LayerTypeEthernet)
	ethernetPacket, _ := ethernetLayer.(*layers.Ethernet)
	fmt.Fprintf(&buf, "%s|", ethernetPacket.SrcMAC)
	fmt.Fprintf(&buf, "%s|", ethernetPacket.DstMAC)
	fmt.Fprintf(&buf, "%s|", ethernetPacket.EthernetType)

	// 网络层
	if ethernetPacket.EthernetType.String() == "IPv4" {
		ipLayer := packet.Layer(layers.LayerTypeIPv4)
		ip, _ := ipLayer.(*layers.IPv4)
		fmt.Fprintf(&buf, "%d|", ip.IHL)
		fmt.Fprintf(&buf, "%d|", ip.TOS)
		fmt.Fprintf(&buf, "%d|", ip.Id)
		fmt.Fprintf(&buf, "%s|", ip.Flags)
		fmt.Fprintf(&buf, "%d|", ip.FragOffset)
		fmt.Fprintf(&buf, "%d|", ip.TTL)
		fmt.Fprintf(&buf, "%s|", ip.SrcIP)
		fmt.Fprintf(&buf, "%s|", ip.DstIP)
		fmt.Fprintf(&buf, "%s|", ip.NextLayerType())
	}

	// 传输层
	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	if tcpLayer != nil {
		tcp, _ := tcpLayer.(*layers.TCP)
		fmt.Fprintf(&buf, "%s|", tcp.SrcPort)
		fmt.Fprintf(&buf, "%s|", tcp.DstPort)
		fmt.Fprintf(&buf, "%d|", tcp.Seq)
		fmt.Fprintf(&buf, "%d|", tcp.Ack)
		fmt.Fprintf(&buf, "%t|", tcp.FIN)
		fmt.Fprintf(&buf, "%t|", tcp.SYN)
		fmt.Fprintf(&buf, "%t|", tcp.RST)
		fmt.Fprintf(&buf, "%t|", tcp.PSH)
		fmt.Fprintf(&buf, "%t|", tcp.ACK)
		fmt.Fprintf(&buf, "%t|", tcp.URG)
		fmt.Fprintf(&buf, "%t|", tcp.ECE)
		fmt.Fprintf(&buf, "%t|", tcp.CWR)
		fmt.Fprintf(&buf, "%t|", tcp.NS)
		fmt.Fprintf(&buf, "%d|", tcp.Window)
		fmt.Fprintf(&buf, "%d", tcp.Checksum)
	}
	return buf.String()

}
